from iou import calculate_iou
class Loss_yolov1(nn.Module):
  def __init__(self):
        super(Loss_yolov1， self).__init__()
  def forward(self， pred， labels):
    """
    pred: (batchsize，30，7，7)的网络输出数据
    labels: (batchsize，30，7，7)的样本标签数据
    :return: 当前批次样本的平均损失
    """
    num_gridx，num_gridy=labels.size()[-2:] # 划分网格数量

    n_batch = labels.size()[0]  # batchsize的大小

    center_pos_loss = 0.    # 中心点坐标损失
    width_height_loss = 0.  # 物体宽高损失
    conf_loss = 0.
    p_class_loss = 0.

    # 可以考虑用矩阵运算进行优化，提高速度，为了准确起见，这里还是用循环
    for batch_i in range(n_batch):  # batchsize循环
     for x_n in range(num_gridx):  # x方向网格循环
      for y_m in range(num_gridy):  # y方向网格循环
        if labels[batch_i，4，x_n，y_m]==1:  
           # 如果包含物体，将数据(px，py，w，h)转换为(x1，y1，x2，y2)
           # 先将px，py转换为cx，cy，
           # 即相对网格的位置转换为标准化后实际的bbox中心位置cx，xy
           # 然后再利用(cx-w/2，cy-h/2，cx+w/2，cy+h/2)转换为xyxy形式，用于计算iou
           # 预测框中心点在区域中的相对位置 转化为 
           # 预测框左上角在整个图像中的相对位置
           cx=(pred[batch_i，0， x_n， y_m]  + y_m) / num_gridx
           w =  pred[batch_i， 2， x_n， y_m]
           cy=(pred[batch_i，1， x_n， y_m] + x_n) / num_gridy
           h = pred[batch_i， 3， x_n， y_m]
           bbox1_pred_xyxy=(cx-w/2，cy-h/2，cx+w/2，cy+h/2)

           cx=(pred[batch_i，5，x_n，y_m] + y_m) / num_gridx
           w = pred[batch_i， 7， x_n， y_m]
           cy=(pred[batch_i，6， x_n，y_m] + x_n) / num_gridy
           h = pred[batch_i， 8， x_n， y_m]
           bbox2_pred_xyxy=(cx-w/2，cy-h/2，cx+w/2，cy+h/2)

           g_cx=(labels[batch_i，0， x_n，y_m]+y_m) / num_gridx
           g_w = labels[batch_i， 2， x_n， y_m]
           g_cy=(labels[batch_i，1， x_n ，y_m]+x_n) / num_gridy
           g_h = labels[batch_i，3， x_n， y_m]
           bbox_gt_xyxy=(g_cx-g_w/2，g_cy-g_h/2，g_cx+g_w/2，g_cy+g_h/2)

          iou1 = calculate_iou(bbox1_pred_xyxy， bbox_gt_xyxy)
          iou2 = calculate_iou(bbox2_pred_xyxy， bbox_gt_xyxy)
          # 选择iou大的bbox作为负责物体
          if iou1 >= iou2:
            # 中心位置损失
            center_pos_loss += 5 *(torch.sum(
                 (pred[batch_i，0:2， x_n，y_m]-labels[batch_i，0:2，x_n，y_m])**2))

            #宽度，高度损失
            width_height_loss +=5*(torch.sum(             
                (labels[batch_i，2， x_n，y_m].sqrt()
                 -pred[batch_i，2， x_n，y_m].sqrt())**2#宽度
                + (labels[batch_i，3， x_n，y_m].sqrt()
                -pred[batch_i，3， x_n，y_m].sqrt())**2)#宽度
                ) 

            #正类置信，负类置信损失
            conf_loss+=(torch.sum(
                    (labels[batch_i，4， x_n， y_m] -pred[batch_i，4， x_n， y_m])**2) 
                     #正类置信损失
                     + 0.5 * torch.sum(
                       (labels[batch_i，9，x_n，y_m] -pred[batch_i，9，x_n，y_m])**2) 
                       #负类置信损失
              ) 

            #分类损失
            p_class_loss += torch.sum(
              (pred[batch_i，10:，x_n，y_m] -labels[batch_i，10:，x_n，y_m])**2  )
          else:
            # 中心位置损失
            center_pos_loss += 5 * (torch.sum(
                     (pred[batch_i， 5:7， x_n， y_m]
                     -labels[batch_i，5:7，x_n，y_m])**2))

            width_height_loss += 5 * (
               torch.sum(
                 (labels[batch_i，7，x_n， y_m].sqrt()
                  - pred[batch_i， 7，x_n，y_m].sqrt())**2# 宽度损失 

                 +(labels[batch_i， 8， x_n， y_m].sqrt()
                  - pred[batch_i，8，x_n，y_m].sqrt() **2 ) # 高度损失 
                 )

            # 置信损失
            conf_loss+=(torch.sum(
                 (labels[batch_i，9，x_n， y_m]
                   -pred[batch_i，9，x_n，y_m])**2) #正类置信损失
                  +0.5*torch.sum(
                    (labels[batch_i，4，x_n， y_m]
                    -pred[batch_i，4，x_n，y_m])**2)#负类置信损失
                ) 

            # 分类损失
            p_class_loss += torch.sum((pred[batch_i， 10:， x_n， y_m]
                                            -labels[batch_i，10:，x_n，y_m])**2 )
        else:  # 如果不包含物体 负类置信损失
   conf_loss += 0.5 * torch.sum((pred[batch_i，[4，9]，x_n ， y_m])**2 ) #负类置信损失 

   all_loss=center_pos_loss + width_height_loss + conf_loss + p_class_loss
   return all_loss / n_batch
